import serial
import struct
import time
import logging
import sys
from serial.tools import list_ports

# 配置日志
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
)
logger = logging.getLogger(__name__)
# ------------------------------- 以下是485通信代码 -------------------------------
def crc8(datagram):
    """计算CRC8校验值，多项式0x07，初始值0x00"""
    crc = 0
    for byte in datagram:
        current_byte = byte
        for _ in range(8):
            bit = (crc >> 7) ^ (current_byte & 0x01)
            if bit:
                crc = (crc << 1) ^ 0x07
            else:
                crc = (crc << 1)
            crc &= 0xFF  # 保持8位
            current_byte >>= 1
    return crc

def build_command_frame(motor_count, ids, command_types, data_list):
    """
    构造主机到电机控制器的数据帧
    :param motor_count: 电机数量（1-8）
    :param ids: 每个电机的ID列表，长度等于motor_count
    :param command_types: 每个电机的指令类型列表
    :param data_list: 每个电机的数据列表，元素为bytes类型
    :return: 完整的字节数据帧
    """
    if motor_count < 1 or motor_count > 8:
        raise ValueError("电机数量必须在1-8之间")
    if len(ids) != motor_count or len(command_types) != motor_count or len(data_list) != motor_count:
        raise ValueError("参数长度与电机数量不匹配")
    
    # 构造所有指令块
    instruction_blocks = []
    for i in range(motor_count):
        id_byte = bytes([ids[i]])
        cmd_byte = bytes([command_types[i]])
        data_bytes = data_list[i]
        block = id_byte + cmd_byte + data_bytes
        instruction_blocks.append(block)
    
    # 合并指令块
    instruction_data = b''.join(instruction_blocks)
    
    # 计算数据长度字段：2(FF FF) + 1（自身） + 1（电机数量） + 指令块总长度 + 1（CRC）
    data_length = 2 + 1 + 1 + len(instruction_data) + 1
    if data_length > 128:
        raise ValueError("数据长度超过最大限制128字节")
        
    # 构造完整数据帧
    frame = b'\xff\xff'  # 字头
    frame += bytes([data_length])
    frame += bytes([motor_count])
    frame += instruction_data
    crc_value = crc8(frame)
    frame += bytes([crc_value])
    
    return frame

def receive_response(ser, expected_length=15, timeout=1):
    """接收响应数据，并校验结构"""
    start_time = time.time()
    buffer = bytearray()
    while time.time() - start_time < timeout:
        if ser.in_waiting > 0:
            buffer += ser.read(ser.in_waiting)
            # 查找字头0xFF 0xFF
            while len(buffer) >= 2:
                pos = buffer.find(b'\xff\xff')
                if pos == -1:
                    # 没有找到字头，保留最后一个字节继续查找
                    buffer = buffer[-1:] if buffer else bytearray()
                    break
                else:
                    # 找到字头，截取后续数据
                    buffer = buffer[pos:]
                    if len(buffer) < expected_length:
                        # 数据不足，继续等待
                        break
                    else:
                        # 提取完整帧
                        frame = buffer[:expected_length]
                        buffer = buffer[expected_length:]
                        return frame
        time.sleep(0.001)
    return None

def parse_stat_response(frame):
    """解析查询指令的响应数据"""
    if len(frame) != 15:
        raise ValueError("响应帧长度必须为15字节")
    # 提取数据部分和CRC
    data_part = frame[0:14]
    received_crc = frame[14]
    # 计算CRC
    calculated_crc = crc8(data_part)
    if calculated_crc != received_crc:
        raise ValueError(f"CRC校验失败: 计算值{calculated_crc:02X}, 接收值{received_crc:02X}")
    # 解析数据字段
    flag = data_part[3]
    if flag != 0:
        raise ValueError(f"标志位不为零, flag={flag}")
    cmd_count = struct.unpack('<H', data_part[4:6])[0]
    trap_status = data_part[6]
    temperature = struct.unpack('b', data_part[7:8])[0]
    pos = struct.unpack('<i', data_part[8:12])[0]
    voltage = struct.unpack('<h', data_part[12:14])[0]
    return [cmd_count, trap_status, temperature, pos, voltage]

def parse_other_response(response_frame):
    if len(response_frame) != 5:
        raise ValueError("响应帧长度必须为5字节")
    logger.debug(f"接收到响应帧: {response_frame.hex()}")
    # 计算CRC
    calculated_crc = crc8(response_frame[0:4])
    received_crc = response_frame[4]
    if calculated_crc != received_crc:
        raise ValueError(f"CRC校验失败: 计算值{calculated_crc:02X}, 接收值{received_crc:02X}")
    # 解析数据字段
    flag = response_frame[3]
    if flag != 0:
        raise ValueError(f"标志位不为零, flag={flag}")
    return True

# 使能、禁用编号为id_list的电机，可同步控制多个
# 用法举例，禁用编号为1,2的电机：cmd_enable(ser,[1,2],0)
def cmd_enable(ser, id_list, en):
    # 构造使能指令（0x01）
    data = bytes([en])  # 启用动力
    motor_count = len(id_list) 
    enable_frame = build_command_frame(
        motor_count, 
        id_list, 
        [0x01] * motor_count, 
        [data] * motor_count)
    logger.debug(f"使能指令数据帧: {enable_frame.hex()}")
    ser.write(enable_frame)

    # 接收响应
    response_frame = receive_response(ser, 5)
    return parse_other_response(response_frame)

def cmd_stat(ser, id):
    # 构造查询指令（0x00）
    motor_count = 1
    ids = [id]
    command_types = [0x00]
    data_list = [bytes()]  # 无数据部分
    stat_frame = build_command_frame(motor_count, ids, command_types, data_list)
    #logger.debug(f"电机编号: {id}, 查询指令数据帧: {stat_frame.hex()}")
    ser.write(stat_frame)

    # 接收并解析响应
    response_frame = receive_response(ser, 15)
    if response_frame:
        #logger.debug(f"接收到响应帧: {response_frame.hex()}")
        parsed_data = parse_stat_response(response_frame)
        # logger.debug(f"[count, trap, temp, pos, voltage]={parsed_data}")
        return parsed_data
    else:
        logger.debug("未接收到响应")
        return None

def cmd_wait_motion(ser, id):
    wait = True
    start_pos = None
    end_pos = None
    while(wait):
        resp = cmd_stat(ser, id)
        pos = resp[3]
        if start_pos == None:
            start_pos = pos
        end_pos = pos
        if resp[1] == 0:
            wait = False
    return (end_pos, end_pos - start_pos)

def cmd_trap(ser, id_list, zero, trap_list):
    count = len(id_list)
    cmd_type = 0x03 if zero else 0x02
    data_bytes = [None]*count
    for i in range (count):
        data_bytes[i] = struct.pack('<i', trap_list[i][0])  # int32 x1
        data_bytes[i] += struct.pack('<h', trap_list[i][1]) # int16 v1
        data_bytes[i] += struct.pack('<h', trap_list[i][2]) # int16 vmax
        data_bytes[i] += struct.pack('<h', trap_list[i][3]) # int16 a
        data_bytes[i] += bytes([trap_list[i][4]])           # uint8 max_current
    print(len(data_bytes[0]))
    trap_frame = build_command_frame(count, id_list, [cmd_type]*count, data_bytes)
    logger.debug(f"梯形运动指令数据帧: {trap_frame.hex()}")
    ser.write(trap_frame)

    # 接收响应
    response_frame = receive_response(ser, 5)
    return parse_other_response(response_frame)

# ------------------------------- 以下是回零点代码 -------------------------------
def cmd_zero(ser):
    # 根据测试情况修改
    arm2_zero = 11117
    arm4_zero = 9336
    current = 15 #电流15%
    # 禁用全部电机4
    cmd_enable(ser, [1,2,3,4], 0)
    # 手指1回零点
    cmd_enable(ser, [1, 2], True)
    finger1_pos, _ = cmd_wait_motion(ser, 1)
    # 位置，终止速度，最大速度30RPM，加速度，电流current
    cmd_trap(ser, [1], True, [[finger1_pos + 8192, 0, 30, 20, current]]) 
    finger1_pos, motion_distance = cmd_wait_motion(ser, 1)
    logger.debug(f"手指1回零点过程移动距离: {motion_distance}, 目前位置: {finger1_pos}")
    if (abs(motion_distance) > 7500):
        logger.debug("移动距离过长，回零点失败")
        return False
    finger1_pos -= 300
    cmd_trap(ser, [1], False, [[finger1_pos, 0, 30, 20, 100]]) # 回退一点
    cmd_wait_motion(ser, 1)
    # 旋转臂2回零点
    arm2_pos, _ = cmd_wait_motion(ser, 2)
    finger1_pos -= arm2_zero - arm2_pos
    cmd_trap(ser, [1, 2], False, [[finger1_pos, 0, 200, 20, 100], [arm2_zero, 0, 200, 20, 100]])
    cmd_wait_motion(ser, 1)
    arm2_pos = arm2_zero

    # 手指3回零点
    cmd_enable(ser, [3, 4], True)
    finger3_pos, _ = cmd_wait_motion(ser, 3)
    # 位置，终止速度，最大速度30RPM，加速度，电流current
    cmd_trap(ser, [3], True, [[finger3_pos + 8192, 0, 30, 20, current]]) 
    finger3_pos, motion_distance = cmd_wait_motion(ser, 3)
    logger.debug(f"手指3回零点过程移动距离: {motion_distance}, 目前位置: {finger3_pos}")
    if (abs(motion_distance) > 7500):
        logger.debug("移动距离过长，回零点失败")
        return False
    finger3_pos -= 300
    cmd_trap(ser, [3], False, [[finger3_pos, 0, 30, 20, 100]]) # 回退一点
    cmd_wait_motion(ser, 3)
    # 旋转臂4回零点
    arm4_pos, _ = cmd_wait_motion(ser, 4)
    finger3_pos -= arm4_zero - arm4_pos
    cmd_trap(ser, [3, 4], False, [[finger3_pos, 0, 200, 20, 100], [arm4_zero, 0, 200, 20, 100]])
    cmd_wait_motion(ser, 3)
    arm4_pos = arm4_zero
    logger.debug(f"finger1_pos={finger1_pos}")
    logger.debug(f"arm2_pos={arm2_pos}")
    logger.debug(f"finger3_pos={finger3_pos}")
    logger.debug(f"arm4_pos={arm4_pos}")
    # cmd_enable(ser, [1,3], 0)
    return (finger1_pos, arm2_pos, finger3_pos, arm4_pos)

# ------------------------------- 以下是运动控制代码 -------------------------------
# // 拧180度
# // twist_cube_180(LEFT); or twist_cube_180(RIGHT);
# void twist_cube_180(int lr)
# {
#     move_arm_block(lr, 1600, cfg.V_START_ARM_L, cfg.V_MAX_ARM_L, cfg.A_MAX_ARM_L);
# }
# // 拧90度
# void twist_cube_90(int lr, int cw_ccw, int skip_back_step)
# {
#     int dir = cw_ccw ? -1 : 1;
#     // 转90度 
#     move_arm_block(lr, dir * 800, cfg.V_START_ARM_L, cfg.V_MAX_ARM_L, cfg.A_MAX_ARM_L);
#     if(!skip_back_step)
#     {
#         // 对侧手指锁紧
#         move_finger_block(!lr, cfg.FINGER_OFFSET_LOCK, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#         sleep_us(cfg.DELAY_US_AFTER_FINGER_LOCK);
#         // 手指松开
#         move_finger_block(lr, cfg.FINGER_OFFSET_MAX, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#         // 转90度
#         move_arm_block(lr, dir * 800, cfg.V_START_ARM, cfg.V_MAX_ARM, cfg.A_MAX_ARM);
#         // 手指缩回
#         move_finger_block(lr, cfg.FINGER_OFFSET_SPIN, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#         // 对侧手指缩回
#         move_finger_block(!lr, cfg.FINGER_OFFSET_SPIN, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     }
# }
# void flip_cube_180(int lr)
# {
#     // 手指锁紧
#     move_finger_block(lr, cfg.FINGER_OFFSET_LOCK, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     sleep_us(cfg.DELAY_US_AFTER_FINGER_LOCK);
#     // 对侧手指缩回
#     move_finger_block(!lr, cfg.FINGER_OFFSET_MAX, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     // 转180
#     move_arm_block(lr, 1600, cfg.V_START_ARM_L, cfg.V_MAX_ARM_L, cfg.A_MAX_ARM_L);
#     // 手指归位
#     move_finger_block(!lr, cfg.FINGER_OFFSET_SPIN, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     move_finger_block(lr, cfg.FINGER_OFFSET_SPIN, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
# }
# void flip_cube_90(int lr, int cw_ccw, int skip_back_step)
# {
#     int dir = cw_ccw ? -1 : 1;
#     // 手指锁紧
#     move_finger_block(lr, cfg.FINGER_OFFSET_LOCK, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     sleep_us(cfg.DELAY_US_AFTER_FINGER_LOCK);
#     // 对侧手指松开
#     move_finger_block(!lr, cfg.FINGER_OFFSET_MAX, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     // 转90
#     move_arm_block(lr, dir * 800, cfg.V_START_ARM_L, cfg.V_MAX_ARM_L, cfg.A_MAX_ARM_L);
#     if(!skip_back_step)
#     {
#         // 手指归位
#         move_finger_block(!lr, cfg.FINGER_OFFSET_SPIN, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#         move_finger_block(!lr, cfg.FINGER_OFFSET_LOCK, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#         sleep_us(cfg.DELAY_US_AFTER_FINGER_LOCK);
#         move_finger_block(lr, cfg.FINGER_OFFSET_MAX, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#         // 转90度
#         move_arm_block(lr, dir * 800, cfg.V_START_ARM, cfg.V_MAX_ARM, cfg.A_MAX_ARM);
#     }
#     // 手指缩回
#     move_finger_block(lr, cfg.FINGER_OFFSET_SPIN, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     // 对侧手指缩回
#     move_finger_block(!lr, cfg.FINGER_OFFSET_SPIN, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
# }
class MotionCtrl:
    def __init__(self, ser, finger1_pos, arm2_pos, finger3_pos, arm4_pos):
        self.finger_zero = [finger1_pos, finger3_pos]
        self.arm_zero = [arm2_pos, arm4_pos]

        self.finger_offset = [0, 0]
        self.arm_offset = [0, 0]

        self.FINGER_CLAMP  = 0      # 手指锁紧
        self.FINGER_INIT   = -950   # 手指初始位置，用于放置魔方
        self.FINGER_MAX    = -6400  # 手指移动到最远位置
        self.MAX_CURRENT   = 100    # 最大电流
        self.CLAMP_CURRENT = 33     # 夹持魔方时的电流
        self.V_SLOW        = 200
        self.V_CLAMP       = 30
        self.A_SLOW        = 100
        self.V_FAST        = 1500
        self.A_FAST        = 400

        self.ser = ser
        pass
    
    def move_two_finger_slow(self, target, current):
        self.finger_offset[0] = target
        self.finger_offset[1] = target
        # 手臂目标位置 = arm_zero + arm_offset
        # 手指目标位置 = finger_zero - arm_offset + finger_offset
        finger1 = self.finger_zero[0] + self.arm_offset[0] + self.finger_offset[0]
        finger3 = self.finger_zero[1] - self.arm_offset[1] + self.finger_offset[1]
        cmd_trap(self.ser, [1, 3], False, 
                 [[finger1, 0, self.V_SLOW, self.A_SLOW, current], 
                  [finger3, 0, self.V_SLOW, self.A_SLOW, current]])
        cmd_wait_motion(self.ser, 3)

    def two_finger_init(self):
        self.move_two_finger_slow(self.FINGER_INIT, self.MAX_CURRENT)

    def two_finger_clamp(self):
        self.move_two_finger_slow(self.FINGER_CLAMP, self.CLAMP_CURRENT)
    
    def two_finger_max(self):
        self.move_two_finger_slow(self.FINGER_MAX, self.MAX_CURRENT)
    
    # 旋转
    def move_arm(self, angle ,left, clamp):
        if left:
            id_list = [3, 4]
            index = 1
        else:
            id_list = [1, 2]
            index = 0

        self.arm_offset[index] += 4096 * (angle // 90)

        # 手臂目标位置 = arm_zero + arm_offset
        arm_target = self.arm_zero[index] + self.arm_offset[index]
        # 存在齿轮，所以需要和手臂电机旋转方向相反，转速绝对值相同，才能保证相对静止
        # 手指目标位置 = finger_zero - arm_offset + finger_offset
        finger_target = self.finger_zero[index] - self.arm_offset[index] + self.finger_offset[index]

        if clamp:
            finger_current = self.CLAMP_CURRENT
            speed = self.V_SLOW
            accel = self.A_SLOW
        else:
            finger_current = self.MAX_CURRENT
            speed = self.V_FAST
            accel = self.A_FAST

        cmd_trap(self.ser, id_list, False, 
                 [[finger_target, 0, speed, accel, finger_current], 
                  [arm_target, 0, speed, accel, self.MAX_CURRENT]])
        cmd_wait_motion(self.ser, id_list[0])

    # 手指伸缩
    def move_finger(self, clamp, left):
        if left:
            id_list = [3]
            index = 1
        else:
            id_list = [1]
            index = 0
        
        if clamp:
            if self.finger_offset[index] == self.FINGER_CLAMP:
                logger.debug("already in clamp status")
                return
            else:
                # 快速接近魔方,接近魔方时的转速保持在V_CLAMP
                self.finger_offset[index] = self.FINGER_INIT
                finger_target = self.finger_zero[index] - self.arm_offset[index] + self.finger_offset[index]
                cmd_trap(self.ser, id_list, False, [[finger_target, self.V_CLAMP, self.V_FAST, self.A_FAST, self.MAX_CURRENT]])
                # 减速，夹持魔方
                self.finger_offset[index] = self.FINGER_CLAMP
                finger_target = self.finger_zero[index] - self.arm_offset[index] + self.finger_offset[index]
                cmd_trap(self.ser, id_list, False, [[finger_target, 0, self.V_CLAMP, self.A_SLOW, self.CLAMP_CURRENT]])
                cmd_wait_motion(self.ser, id_list[0])
                
        else:
            if self.finger_offset[index] == self.FINGER_MAX:
                logger.debug("already in unclamp status")
                return
            else:
                # 快速松开魔方
                self.finger_offset[index] = self.FINGER_MAX
                finger_target = self.finger_zero[index] - self.arm_offset[index] + self.finger_offset[index]
                cmd_trap(self.ser, id_list, False, [[finger_target, 0, self.V_FAST, self.A_FAST, self.MAX_CURRENT]])
                cmd_wait_motion(self.ser, id_list[0])
            

    def flip_cube_180(self, left):
        # 对侧手指缩回
        self.move_finger(False, not left)
        # 转180
        self.move_arm(180, left, True)
        # 手指归位
        self.move_finger(True, not left)

    # void flip_cube_180(int lr)
# {
#     // 手指锁紧
#     move_finger_block(lr, cfg.FINGER_OFFSET_LOCK, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     sleep_us(cfg.DELAY_US_AFTER_FINGER_LOCK);
#     // 对侧手指缩回
#     move_finger_block(!lr, cfg.FINGER_OFFSET_MAX, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     // 转180
#     move_arm_block(lr, 1600, cfg.V_START_ARM_L, cfg.V_MAX_ARM_L, cfg.A_MAX_ARM_L);
#     // 手指归位
#     move_finger_block(!lr, cfg.FINGER_OFFSET_SPIN, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
#     move_finger_block(lr, cfg.FINGER_OFFSET_SPIN, cfg.V_START_FINGER, cfg.V_MAX_FINGER, cfg.A_MAX_FINGER);
# }

# ------------------------------- 以下是测试程序 -------------------------------
def list_serial_ports():
    """列出可用串口"""
    ports = list_ports.comports()
    if not ports:
        print("未检测到可用串口设备")
        return
    print("可用串口设备:")
    for port in ports:
        print(f"  {port.device} - {port.description}")

if __name__ == "__main__":
    if len(sys.argv) > 1:
        serial_port = sys.argv[1]
    else:
        serial_port = '/dev/cu.usbserial-130'  # 默认值

    baud_rate = 1000000  # 1Mbps
    mc = None
    try:
        with serial.Serial(serial_port, baudrate=baud_rate, bytesize=serial.EIGHTBITS,
                        parity=serial.PARITY_NONE, stopbits=serial.STOPBITS_ONE) as ser:
            print("\n串口连接成功，进入交互模式")
            while True:
                print("\n请选择测试项目:")
                print("[1]: 使能全部电机")
                print("[2]: 禁用全部电机")
                print("[3]: 查询电机角度")
                print("[4]: 回零点")
                print("[5]: 测试拧魔方")
                print("[q]: 退出程序")
                
                choice = input("请输入选项(1/2/q/...) >> ").strip().lower()
                
                if choice in ['exit', 'quit', 'q']:
                    print("退出程序...")
                    break
                elif choice == '1':
                    success = cmd_enable(ser, [1,2,3,4], 1)
                    print("执行结果:", "成功" if success else "失败")
                elif choice == '2':
                    success = cmd_enable(ser, [1,2,3,4], 0)
                    print("执行结果:", "成功" if success else "失败")
                elif choice == '3':
                    exit = False
                    for i in range(100):
                        pos = [None] * 4
                        for id in (1,2,3,4):
                            resp = cmd_stat(ser, id)
                            if resp == None:
                                print(f"未收到控制器回复，控制器编号={id}")
                                exit = True
                            else:
                                pos[id-1] = resp[3]
                        print(f"当前电机角度: {pos}")
                        if exit:
                            break
                elif choice == '4':
                    zero = cmd_zero(ser)
                    mc = MotionCtrl(ser, zero[0], zero[1], zero[2], zero[3])
                elif choice == '5':
                    mc.two_finger_init()
                elif choice == '6':
                    mc.two_finger_clamp()
                elif choice == '7':
                    mc.two_finger_max()
                elif choice == '8':
                    mc.flip_cube_180(True)
                else:
                    print("无效选项，请重新输入")


    except serial.SerialException as e:
        logger.error(f"打开串口 {serial_port} 失败: {str(e)}")
        list_serial_ports()
        sys.exit(1)
    except KeyboardInterrupt:
        logger.info("程序被用户中断")

